import math, re, sys

g_RARE = "_RARE_"

def read_counts_file(filename):
    counts = {}
    trigram_counts = {}
    f = open(filename)
    for line in f:
        tokens = line.split()
        if len(tokens) < 4:
            continue
        if tokens[1] == "WORDTAG":
            count = int(tokens[0])
            tag = tokens[2]
            word = tokens[3]
            counts[(tag, word)] = count
        elif tokens[1] == "3-GRAM":
            count = int(tokens[0])
            trigram_counts[tuple(tokens[2:5])] = count
    f.close()
    return counts, trigram_counts

def compute_prob(counts):
    counts = collapse_rare_words(counts)
    prob = {}
    tag_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        accum(tag_counts, tag, count)
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        tag_count = tag_counts[tag]
        dict2_add(prob, word, tag, count * 1.0 / tag_count)
    return prob

def compute_trigram_prob(counts):
    cond_counts = {}
    prob = {}
    for pair, count in counts.viewitems():
        cond = pair[0:-1]
        accum(cond_counts, cond, count)
    for pair, count in counts.viewitems():
        cond= pair[0:-1]
        var = pair[-1]
        cond_count = cond_counts[cond]
        dict2_add(prob, cond, var, count * 1.0 / cond_count)
    return prob

def collapse_rare_words(counts):
    word_counts = {}
    for pair, count in counts.viewitems():
        word = pair[1]
        accum(word_counts, word, count)
    rare_words = set([word for word, count in word_counts.viewitems() if count < 5])
    new_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        if word in rare_words:
            word = decide_rare_word(word)
        accum(new_counts, (tag, word), count)
    return new_counts

rare_names = {}
rare_names['_PROPER_'] = re.compile('^[A-Z]+$')
rare_names['_ABBR_'] = re.compile('^[A-Z\.]+$')
rare_names['_NUMERAL_'] =re.compile('^[\d\-\.]+$')

def decide_rare_word(w):
    for name, pattern in rare_names.viewitems():
        if pattern.match(w) <> None:
            return name
    return g_RARE

def accum(d, key, value):
    d[key] = d.get(key, 0) + value

def dict2_add(d, key1, key2, value):
    d2 = d.get(key1, {})
    d2[key2] = value
    d[key1] = d2

def tag_file(sent_file, tagger):
    f = open(sent_file)
    sent = []
    for line in f:
        tokens = line.split()
        if len(tokens) <> 1:
            tagger(sent)
            sent = []
            continue
        word = tokens[0]
        sent += [word]
    f.close

def get_states(trans_prob):
    return unique([s for s, s2 in trans_prob.viewkeys() if s <> "STOP" and s <> "*"])

def unique(ls):
    return [x for x in set(ls)]
    
def get_words(sent):
    words = [w for w in set(sent)]
    word_idx = dict([(w, i) for i, w in enumerate(words)])
    return words, word_idx

def replace_by_rare(sent, prob):
    return [w if w in prob else decide_rare_word(w) for w in sent]

def init_list(m, n):
    return [[0] * n for x in range(m)]

def init_list3(m, n, r):
    return [[[0] * r for y in range(n)] for x in range(m)]

def argmax_max(ls):
    return max(enumerate(ls), key = lambda x : x[1])

def trigram_tag(sent, prob, trans_prob):
    states = get_states(trans_prob)
    states += '*'
    N = len(states)
    star_idx = N - 1
    origin_sent = sent
    sent = replace_by_rare(sent, prob)
    words, word_idx = get_words(sent)
    sent = [word_idx[w] for w in sent]
    W = len(words)
    prob_b = init_list(N, W)
    for w in range(W):
        p = prob.get(words[w], {})
        for s in range(N):
            prob_b[s][w] = p.get(states[s], 0)
    def lookup(s1, s2, s3):
        return trans_prob.get((s1, s2), {}).get(s3, 0)
    prob_a = init_list3(N, N, N)
    for s1 in range(N):
        for s2 in range(N):
            for s3 in range(N):
                prob_a[s1][s2][s3] = lookup(states[s1], states[s2], states[s3])
    T = len(sent)
    v = init_list3(T+1, N, N)
    bv = init_list3(T+1, N, N)
    v[0][star_idx][star_idx] = 1
    for t in range(1, T+1):
        for s2 in range(N):
            for s3 in range(N):
                bv[t][s2][s3], v[t][s2][s3] = argmax_max([v[t-1][s1][s2] * prob_a[s1][s2][s3] for s1 in range(N)])
                v[t][s2][s3] *= prob_b[s3][sent[t-1]]
    bvf, vf = max([((s1,s2), v[T][s1][s2] * lookup(states[s1], states[s2], "STOP")) for s1 in range(N) for s2 in range(N)], key = lambda x : x[1])
    tags = [0] * T
    scores = [0.0] * T
    tags[T - 1] = bvf[1]
    tags[T - 2] = bvf[0]
    scores[T - 1] = v[T][tags[T - 2]][tags[T - 1]]
    for t in range(T-3, -1, -1):
        tags[t] = bv[t+3][tags[t+1]][tags[t+2]]
        scores[t+1] = v[t+2][tags[t]][tags[t+1]]
    scores[0] = v[1][star_idx][tags[0]]
    tags = [states[s] for s in tags];
    for t in range(T):
        print origin_sent[t], tags[t], math.log(scores[t]+1e-10, 2)
    print

counts_file = "ner.counts"
sent_file = "ner_dev.dat"

if len(sys.argv) >= 3:
    counts_file = sys.argv[1]
    sent_file = sys.argv[2]

counts, trigram_counts = read_counts_file(counts_file)
# print bigram_counts
prob = compute_prob(counts)
trans_prob = compute_trigram_prob(trigram_counts)
# print prob
tag_file(sent_file, lambda s : trigram_tag(s, prob, trans_prob))
